﻿using BuddHa.Common.SMUtil.Implementation.V1;
using Org.BouncyCastle.Crypto;
using Org.BouncyCastle.Crypto.Digests;
using Org.BouncyCastle.Crypto.Parameters;
using Org.BouncyCastle.Math;
using Org.BouncyCastle.Math.EC;
using Org.BouncyCastle.Math.EC.Multiplier;
using Org.BouncyCastle.Security;
using Org.BouncyCastle.Utilities;
using System;

namespace BuddHa.Common.SMUtil.Implementation.V2
{
    internal class SM2Encryptor
    {
        private readonly IDigest _mDigest;

        private ECKeyParameters _mECKey;
        private ECDomainParameters _mECParams;
        private SecureRandom _mRandom;

        public SM2Encryptor()
           : this(new SM3Digest())
        {
        }

        public SM2Encryptor(IDigest digest)
        {
            _mDigest = digest;
        }

        public virtual void Init(ICipherParameters param)
        {

            ParametersWithRandom rParam = (ParametersWithRandom)param;

            _mECKey = (ECKeyParameters)rParam.Parameters;
            _mECParams = _mECKey.Parameters;

            var s = ((ECPublicKeyParameters)_mECKey).Q.Multiply(_mECParams.H);
            if (s.IsInfinity)
            {
                throw new ArgumentException("invalid key: [h]Q at infinity");
            }

            _mRandom = rParam.Random;
        }

        public byte[] Encrypt(byte[] data)
        {
            return Encrypt(data, 0, data.Length);
        }

        private byte[] Encrypt(byte[] input, int inOff, int inLen)
        {

            _mDigest.Reset();

            byte[] c2 = new byte[inLen];

            Array.Copy(input, inOff, c2, 0, c2.Length);

            var multiplier = new FixedPointCombMultiplier();

            byte[] c1;
            ECPoint kPB;
            do
            {
                BigInteger k = NextK();

                ECPoint c1P = multiplier.Multiply(_mECParams.G, k).Normalize();

                c1 = c1P.GetEncoded(false);

                kPB = ((ECPublicKeyParameters)_mECKey).Q.Multiply(k).Normalize();

                Kdf(_mDigest, kPB, c2);
            }
            while (NotEncrypted(c2, input, inOff));

            AddFieldElement(_mDigest, kPB.AffineXCoord);
            _mDigest.BlockUpdate(input, inOff, inLen);
            AddFieldElement(_mDigest, kPB.AffineYCoord);

            byte[] c3 = DigestUtilities.DoFinal(_mDigest);

            return Arrays.ConcatenateAll(c1, c2, c3);
        }

        private bool NotEncrypted(byte[] encData, byte[] input, int inOff)
        {
            for (int i = 0; i != encData.Length; i++)
            {
                if (encData[i] != input[inOff + i])
                {
                    return false;
                }
            }

            return true;
        }

        private BigInteger NextK()
        {
            int qBitLength = _mECParams.N.BitLength;

            BigInteger k;
            do
            {
                k = new BigInteger(qBitLength, _mRandom);
            }
            while (k.SignValue == 0 || k.CompareTo(_mECParams.N) >= 0);

            return k;
        }

        private void AddFieldElement(IDigest digest, ECFieldElement v)
        {
            byte[] p = v.GetEncoded();
            digest.BlockUpdate(p, 0, p.Length);
        }

        private void Kdf(IDigest digest, ECPoint c1, byte[] encData)
        {
            int digestSize = digest.GetDigestSize();
            byte[] buf = new byte[System.Math.Max(4, digestSize)];
            int off = 0;

            IMemoable memo = digest as IMemoable;
            IMemoable copy = null;

            if (memo != null)
            {
                AddFieldElement(digest, c1.AffineXCoord);
                AddFieldElement(digest, c1.AffineYCoord);
                copy = memo.Copy();
            }

            uint ct = 0;

            while (off < encData.Length)
            {
                if (memo != null)
                {
                    memo.Reset(copy);
                }
                else
                {
                    AddFieldElement(digest, c1.AffineXCoord);
                    AddFieldElement(digest, c1.AffineYCoord);
                }

                Pack.UInt32_To_BE(++ct, buf, 0);
                digest.BlockUpdate(buf, 0, 4);
                digest.DoFinal(buf, 0);

                int xorLen = System.Math.Min(digestSize, encData.Length - off);
                Xor(encData, buf, off, xorLen);
                off += xorLen;
            }
        }

        private void Xor(byte[] data, byte[] kdfOut, int dOff, int dRemaining)
        {
            for (int i = 0; i != dRemaining; i++)
            {
                data[dOff + i] ^= kdfOut[i];
            }
        }
    }
}
